#include <iostream>
#include <thread>
#include <atomic>
#include <vector>
#include <cstring>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>

class TCPUDPForwarder {
private:
    int tcp_listen_sock;
    int udp_sock;
    int tcp_client_sock;
    std::atomic<bool> running;
    std::thread tcp_listen_thread;
    std::thread udp_receive_thread;
    std::thread tcp_receive_thread;
    
    struct sockaddr_in udp_target_addr;
    int porta, portb, portc;
    std::string ipu;

public:
    TCPUDPForwarder(int porta, const std::string& ipu, int portb, int portc) 
        : porta(porta), ipu(ipu), portb(portb), portc(portc), 
          tcp_listen_sock(-1), udp_sock(-1), tcp_client_sock(-1), running(false) {
        
        // 设置UDP目标地址
        memset(&udp_target_addr, 0, sizeof(udp_target_addr));
        udp_target_addr.sin_family = AF_INET;
        udp_target_addr.sin_port = htons(portb);
        inet_pton(AF_INET, ipu.c_str(), &udp_target_addr.sin_addr);
    }

    ~TCPUDPForwarder() {
        stop();
    }

    bool start() {
        if (!createTCPSocket() || !createUDPSocket()) {
            return false;
        }

        running = true;
        
        // 启动TCP监听线程
        tcp_listen_thread = std::thread(&TCPUDPForwarder::tcpListenLoop, this);
        
        // 启动UDP接收线程
        udp_receive_thread = std::thread(&TCPUDPForwarder::udpReceiveLoop, this);

        std::cout << "Forwarder started: TCP:" << porta << " <-> UDP:" << portc 
                  << " -> " << ipu << ":" << portb << std::endl;
        
        return true;
    }

    void stop() {
        running = false;
        
        // 关闭socket以唤醒阻塞的线程
        if (tcp_listen_sock != -1) {
            close(tcp_listen_sock);
            tcp_listen_sock = -1;
        }
        if (udp_sock != -1) {
            close(udp_sock);
            udp_sock = -1;
        }
        if (tcp_client_sock != -1) {
            close(tcp_client_sock);
            tcp_client_sock = -1;
        }

        // 等待线程结束
        if (tcp_listen_thread.joinable()) {
            tcp_listen_thread.join();
        }
        if (udp_receive_thread.joinable()) {
            udp_receive_thread.join();
        }
        if (tcp_receive_thread.joinable()) {
            tcp_receive_thread.join();
        }
    }

private:
    bool createTCPSocket() {
        tcp_listen_sock = socket(AF_INET, SOCK_STREAM, 0);
        if (tcp_listen_sock < 0) {
            std::cerr << "Failed to create TCP socket: " << strerror(errno) << std::endl;
            return false;
        }

        // 设置SO_REUSEADDR
        int opt = 1;
        if (setsockopt(tcp_listen_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
            std::cerr << "Failed to set TCP socket options: " << strerror(errno) << std::endl;
            close(tcp_listen_sock);
            return false;
        }

        struct sockaddr_in addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_addr.s_addr = INADDR_ANY;
        addr.sin_port = htons(porta);

        if (bind(tcp_listen_sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
            std::cerr << "Failed to bind TCP socket: " << strerror(errno) << std::endl;
            close(tcp_listen_sock);
            return false;
        }

        if (listen(tcp_listen_sock, 5) < 0) {
            std::cerr << "Failed to listen on TCP socket: " << strerror(errno) << std::endl;
            close(tcp_listen_sock);
            return false;
        }

        return true;
    }

    bool createUDPSocket() {
        udp_sock = socket(AF_INET, SOCK_DGRAM, 0);
        if (udp_sock < 0) {
            std::cerr << "Failed to create UDP socket: " << strerror(errno) << std::endl;
            return false;
        }

        // 绑定到portc用于接收UDP数据
        struct sockaddr_in udp_listen_addr;
        memset(&udp_listen_addr, 0, sizeof(udp_listen_addr));
        udp_listen_addr.sin_family = AF_INET;
        udp_listen_addr.sin_addr.s_addr = INADDR_ANY;
        udp_listen_addr.sin_port = htons(portc);

        if (bind(udp_sock, (struct sockaddr*)&udp_listen_addr, sizeof(udp_listen_addr)) < 0) {
            std::cerr << "Failed to bind UDP socket: " << strerror(errno) << std::endl;
            close(udp_sock);
            return false;
        }

        // 设置socket为非阻塞
        int flags = fcntl(udp_sock, F_GETFL, 0);
        fcntl(udp_sock, F_SETFL, flags | O_NONBLOCK);

        return true;
    }

    void tcpListenLoop() {
        while (running) {
            struct sockaddr_in client_addr;
            socklen_t client_len = sizeof(client_addr);
            
            int new_client = accept(tcp_listen_sock, (struct sockaddr*)&client_addr, &client_len);
            if (new_client < 0) {
                if (errno != EAGAIN && errno != EWOULDBLOCK) {
                    std::cerr << "Accept failed: " << strerror(errno) << std::endl;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                continue;
            }

            // 关闭之前的客户端连接
            if (tcp_client_sock != -1) {
                close(tcp_client_sock);
            }

            tcp_client_sock = new_client;
            char client_ip[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
            std::cout << "TCP client connected: " << client_ip << ":" << ntohs(client_addr.sin_port) << std::endl;

            // 启动TCP数据接收线程
            if (tcp_receive_thread.joinable()) {
                tcp_receive_thread.join();
            }
            tcp_receive_thread = std::thread(&TCPUDPForwarder::tcpReceiveLoop, this);
        }
    }

    void tcpReceiveLoop() {
        std::vector<char> buffer(4096);
        
        while (running && tcp_client_sock != -1) {
            ssize_t bytes_received = recv(tcp_client_sock, buffer.data(), buffer.size(), 0);
            if (bytes_received > 0) {
                // 将TCP数据通过UDP发送到目标地址
                ssize_t bytes_sent = sendto(udp_sock, buffer.data(), bytes_received, 0,
                                           (struct sockaddr*)&udp_target_addr, sizeof(udp_target_addr));
                if (bytes_sent < 0) {
                    std::cerr << "Failed to send UDP data: " << strerror(errno) << std::endl;
                } else {
                    std::cout << "Forwarded " << bytes_received << " bytes from TCP to UDP" << std::endl;
                }
            } else if (bytes_received == 0) {
                std::cout << "TCP client disconnected" << std::endl;
                close(tcp_client_sock);
                tcp_client_sock = -1;
                break;
            } else {
                if (errno != EAGAIN && errno != EWOULDBLOCK) {
                    std::cerr << "TCP receive error: " << strerror(errno) << std::endl;
                    close(tcp_client_sock);
                    tcp_client_sock = -1;
                    break;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            }
        }
    }

    void udpReceiveLoop() {
        std::vector<char> buffer(4096);
        struct sockaddr_in src_addr;
        socklen_t addr_len = sizeof(src_addr);

        while (running && udp_sock != -1) {
            ssize_t bytes_received = recvfrom(udp_sock, buffer.data(), buffer.size(), 0,
                                            (struct sockaddr*)&src_addr, &addr_len);
            
            if (bytes_received > 0) {
                // 将UDP数据转发给TCP客户端
                if (tcp_client_sock != -1) {
                    ssize_t bytes_sent = send(tcp_client_sock, buffer.data(), bytes_received, 0);
                    if (bytes_sent < 0) {
                        std::cerr << "Failed to send TCP data: " << strerror(errno) << std::endl;
                    } else {
                        char src_ip[INET_ADDRSTRLEN];
                        inet_ntop(AF_INET, &src_addr.sin_addr, src_ip, INET_ADDRSTRLEN);
                        std::cout << "Forwarded " << bytes_received << " bytes from UDP:" 
                                  << src_ip << ":" << ntohs(src_addr.sin_port) << " to TCP" << std::endl;
                    }
                } else {
                    std::cout << "Received UDP data but no TCP client connected" << std::endl;
                }
            } else if (bytes_received < 0) {
                if (errno != EAGAIN && errno != EWOULDBLOCK) {
                    std::cerr << "UDP receive error: " << strerror(errno) << std::endl;
                    break;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            }
        }
    }
};

// 接口函数
extern "C" {
    void* create_forwarder(int porta, const char* ipu, int portb, int portc) {
        return new TCPUDPForwarder(porta, ipu, portb, portc);
    }
    
    int start_forwarder(void* forwarder) {
        TCPUDPForwarder* fwd = static_cast<TCPUDPForwarder*>(forwarder);
        return fwd->start() ? 1 : 0;
    }
    
    void stop_forwarder(void* forwarder) {
        TCPUDPForwarder* fwd = static_cast<TCPUDPForwarder*>(forwarder);
        fwd->stop();
    }
    
    void destroy_forwarder(void* forwarder) {
        TCPUDPForwarder* fwd = static_cast<TCPUDPForwarder*>(forwarder);
        delete fwd;
    }
}

// 使用示例
int main() {
    // 创建转发器：TCP监听8080，UDP发送到127.0.0.1:9090，UDP接收在9091
    void* forwarder = create_forwarder(8080, "127.0.0.1", 9090, 9091);
    
    if (start_forwarder(forwarder)) {
        std::cout << "Forwarder started successfully. Press Enter to stop..." << std::endl;
        std::cin.get();
        
        stop_forwarder(forwarder);
    } else {
        std::cerr << "Failed to start forwarder" << std::endl;
    }
    
    destroy_forwarder(forwarder);
    return 0;
}